#include <gmp.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>

typedef struct {
	int bits;
	mpz_t n;  // n
	mpz_t hs; // hs
	mpz_t n2; // n^2
} paillier_pubkey;

typedef struct {
	int bits;
	mpz_t lambda;  // lambda
	mpz_t mu;  // mu
} paillier_prikey;

typedef struct {
	paillier_pubkey* pubkey;
	paillier_prikey* prikey;
	int bits;
} paillier_keypair;


// 生成介于[0, n) 范围内的随机整数
int generate_random_bn(mpz_t rop, mpz_t  n) {	
	gmp_randstate_t state;

	// 初始化随机数生成器状??
	gmp_randinit_default(state);
	mpz_init(rop);
	mpz_urandomm(rop, state, n);
	gmp_randclear(state);

	return 0;
}

// 生成介于 0 ?? 2^k（k 为正整数）之间的随机整数
int generate_random2_bn(mpz_t rop, mp_bitcnt_t  k) {
	gmp_randstate_t state;
	gmp_randinit_default(state);
	mpz_init(rop);
	mpz_urandomb(rop, state, k);
	gmp_randclear(state);
	return 0;
}

// new paillier_keypair
paillier_keypair* paillier_keypair_new() {
	paillier_keypair* keypair =
		(paillier_keypair*)malloc(sizeof(paillier_keypair));
	if (keypair) {
		keypair->bits = 0;
		keypair->prikey = (paillier_prikey*)malloc(sizeof(paillier_prikey));
		if (keypair->prikey) {
			keypair->prikey->bits = 0;
			mpz_init(keypair->prikey->lambda);
			mpz_init(keypair->prikey->mu);
		}
		keypair->pubkey = (paillier_pubkey*)malloc(sizeof(paillier_pubkey));
		if (keypair->pubkey) {
			keypair->pubkey->bits = 0;
			mpz_init(keypair->pubkey->hs);
			mpz_init(keypair->pubkey->n);
			mpz_init(keypair->pubkey->n2);
		}
	}
	return keypair;
}

// free paillier_keypair
void paillier_keypair_free(paillier_keypair* keypair) {
	free(keypair->prikey);
	free(keypair->pubkey);
	free(keypair);
	return;
}


// 生成密钥??
int generate_keypair(paillier_keypair* pk, int bits) {
	pk->bits = bits;
	pk->pubkey->bits = bits;
	pk->prikey->bits = bits;
	mpz_t  times;
	mpz_t  p, q;
	mpz_init(p);
	mpz_init(q);
	
	mpz_t ps1, qs1, temp, three, x, h;
	mpz_init(temp);
	mpz_init(ps1);
	mpz_init(qs1);
	mpz_init(x);
	mpz_init(h);

	gmp_randstate_t state;
	gmp_randinit_default(state);

	// 生成p, q
	mpz_init_set_ui(times, 4);//u_int，times=4
	mpz_urandomb(p, state, bits / 2);//随机数P，0~bits
	mpz_nextprime(p, p);//p设置为p临近大质数
	mpz_mod(temp, p, times); // temp=p%4
	mpz_init_set_ui(three, 3);//u_int，three=3
	
	for (; mpz_cmp(temp, three) != 0;) { // p%4=3
		mpz_urandomb(p, state, bits / 2);//生成0~2^n-1随机数
		mpz_nextprime(p, p);
		mpz_mod(temp, p, times); // p%4
	}
	mpz_init_set(q, p);
	
	int q_p = mpz_cmp(q, p);
	for (; q_p == 0;) {
		mpz_urandomb(q, state, bits / 2);
		mpz_nextprime(q, q);
		mpz_mod(temp, q, times); // q%4
		for (; mpz_cmp(temp, three) != 0;) {            // q%4=3
			mpz_urandomb(q, state, bits / 2);
			mpz_nextprime(q, q);
			mpz_mod(temp, q, times);
		}
		q_p = mpz_cmp(q, p);
	}

	gmp_printf("debug: p: %Zd\t q: %Zd\n", p, q);

	mpz_mul(pk->pubkey->n, p, q); // n=p*q
	mpz_set_ui(times, 1);


	mpz_sub(ps1, p, times);     // ps1 = (p - 1)
	//gmp_printf("\n\ndebug: p: %Zd\t ps1: %Zd", p, ps1);
	mpz_mul(temp, ps1, q); // (p - 1) * q

	mpz_sub(qs1, q, times);     // qs1 = (q - 1)
	mpz_mul(temp, qs1, p); // (q - 1) * p

	mpz_mul(temp, ps1, qs1); // lambda = (p - 1)(q - 1)
	mpz_tdiv_q_2exp(pk->prikey->lambda, temp, 1);  // lambda = lambda除以2
	mpz_invert(pk->prikey->mu, pk->prikey->lambda, pk->pubkey->n); // mu = lambda^-1 mod n

	if (q_p > 0) {
		mpz_urandomm(x, state, p); 
		mpz_add_ui(x, x, 2);	// 生成(2, p)的整??
	}
	else
	{
		mpz_urandomm(x, state, q);
		mpz_add_ui(x, x, 2);	// 生成(2, q)的整??
	}
	
	mpz_set_ui(times, 2);
	mpz_powm(h, x, times, pk->pubkey->n);       // h = x^2 mod n
	//gmp_printf("\n\ndebug: h: %Zd", h);
	mpz_neg(h, h);	// h = - x^2 mod n
	mpz_pow_ui(pk->pubkey->n2, pk->pubkey->n, 2); // n^2
	mpz_powm(pk->pubkey->hs, h, pk->pubkey->n, pk->pubkey->n2); // hs = h^n mod n^2
	//gmp_printf("\n\ndebug: hs: %Zd", pk->pubkey->hs);

	mpz_clear(h);
	mpz_clear(x);
	mpz_clear(temp);
	mpz_clear(ps1);
	mpz_clear(qs1);
	mpz_clear(q);
	mpz_clear(p);
	mpz_clear(times);
	return 0;
}


//m=明文，out=密文,alpha=随机数
int encrypt(mpz_t out, paillier_pubkey* pubkey, mpz_t m) {
	if (pubkey->bits == 0 || (pubkey->bits & 1) != 0) {
		return 1;
	}
	mpz_t temp, temp1, nm, one, alpha;
	mpz_inits(temp, temp1, nm, one, alpha, NULL);
	
	generate_random2_bn(alpha,pubkey->bits / 2);

	mpz_mul(nm, pubkey->n, m);                        // nm = n * m
	mpz_add(temp1, nm, one);                               // n * m + 1
	mpz_powm(temp, pubkey->hs, alpha, pubkey->n2); // hs^alpha mod n^2
	mpz_mul(one, temp1, temp);	// one = (n * m + 1) * (hs^alpha mod n^2)
	mpz_mod(out, one, pubkey->n2); // c = (n * m + 1)hs^alphq mod n^2
	//错误点2，多模了一次
	mpz_clears(alpha, one, nm, temp1, temp, NULL);
	return 0;
}

int decrypt(mpz_t out, paillier_keypair* pk, mpz_t c) {
	mpz_t temp, n2, lnx, temp2;
	mpz_inits(temp, n2, lnx,temp2, NULL);


	mpz_pow_ui(n2, pk->pubkey->n, 2); // n^2
	//mpz_mul(n2, pk->pubkey->n, pk->pubkey->n);
	mpz_powm(lnx, c, pk->prikey->lambda, n2);  // lnx = c^lambda mod n^2
	//gmp_printf("debug: c: %Zd\t lambda: %Zd\t n2: %Zd\t lnx: %Zd\n", c,
	//	pk->prikey->lambda, pk->pubkey->n2, lnx);
	//gmp_printf("debug: c: %d\t lambda: %d\t n2: %d\t lnx: %d\n", c->_mp_size,
	//	pk->prikey->lambda->_mp_size, pk->pubkey->n2->_mp_size, lnx->_mp_size);
	mpz_sub_ui(temp, lnx, 1);	// lnx - 1
	mpz_div(temp2, temp, pk->pubkey->n); // (lnx - 1) / n
	mpz_mul(temp, temp2, pk->prikey->mu);	// (lnx - 1) / n * mu
	mpz_mod(out, temp, pk->pubkey->n);

	mpz_clears(temp2, lnx, n2, temp, NULL);

	return 0;
}

int main() {
	paillier_keypair* keypair = paillier_keypair_new();
	//generate_p_q(keypair, 1024);
	generate_keypair(keypair, 16);

	gmp_printf("\nn:%Zd\t  hs:%Zd\t  n2:%Zd\t  lambda:%Zd\t  mu:%Zd\n",
		keypair->pubkey->n, keypair->pubkey->hs,
		keypair->pubkey->n2, keypair->prikey->lambda,
		keypair->prikey->mu);
	
	mpz_t input, cipher_text, plain_text;
	mpz_inits(input, cipher_text, plain_text, NULL);
	mpz_set_str(input, "1234", 10);
	gmp_printf("明文:%Zd\n", input);
	encrypt(cipher_text, keypair->pubkey, input);
	gmp_printf("密文:%Zd\n", cipher_text);
	decrypt(plain_text, keypair, cipher_text);
	gmp_printf("解密文:%Zd\n", plain_text);
	mpz_clears(plain_text, cipher_text, input, NULL);
	paillier_keypair_free(keypair);
}